feat(audio-io): make IO stages agent-ready (contract + serialization-safe sink) - #2171
feat(audio-io): make IO stages agent-ready (contract + serialization-safe sink)#2171shubhamNvidia wants to merge 3 commits into
Conversation
…idency resolver Shared base imported by the audio stage modules to make them agent-ready: - _agent_ready.py: AgentReady mixin, StageContract/IOSpec/Gates/Role - _residency.py: input residency resolver (file/waveform/auto) - common.py: agent-ready helpers (ensure_mono/ensure_waveform_2d) Excludes the agent orchestration layer (registry/catalog/conformance/planning/agent).
Greptile SummaryThis PR makes the audio IO stages agent-ready by adding a
Confidence Score: 4/5Safe to merge with one fix: the temp-file leak in The new
Important Files Changed
Class Diagram%%{init: {'theme': 'neutral'}}%%
classDiagram
class AgentReady {
+ClassVar AGENT_STATIC: StaticHints | None
+ClassVar BATCH_ONLY: bool
+ClassVar KEY_ROLE_OVERRIDES: Mapping
+describe() StageContract
+describe_static()$ StageContract
}
class StageContract {
+reads: IOSpec
+writes: IOSpec
+reads_one_of: list[IOSpec]
+cardinality: Cardinality
+gates: Gates
+batch_only: bool
+dispatch: Dispatch
+to_dict() dict
}
class Gates {
+writes_to_disk: bool
+requires_gpu: bool
+requires_serializable_input: bool
+sanitizes_output: bool
}
class IOSpec {
+data_keys: list[str]
+segment_data_keys: list[str]
+accepts: list[AudioForm]
+produces: list[ProducedForm]
}
class AudioToDocumentStage {
+BATCH_ONLY = True
+keep_keys: list[str] | None
+drop_keys: tuple[str, ...]
+serialize_segments: bool
+segments_key: str
+describe() StageContract
+_sanitize(data) dict
+_sanitize_nested(value) object
+process_batch(tasks) list
}
class SegmentExtractionStage {
+BATCH_ONLY = True
+output_dir: str
+output_key: str
+describe() StageContract
+process_batch(tasks) list
}
AgentReady <|-- AudioToDocumentStage
AgentReady <|-- SegmentExtractionStage
StageContract *-- Gates
StageContract *-- IOSpec
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
classDiagram
class AgentReady {
+ClassVar AGENT_STATIC: StaticHints | None
+ClassVar BATCH_ONLY: bool
+ClassVar KEY_ROLE_OVERRIDES: Mapping
+describe() StageContract
+describe_static()$ StageContract
}
class StageContract {
+reads: IOSpec
+writes: IOSpec
+reads_one_of: list[IOSpec]
+cardinality: Cardinality
+gates: Gates
+batch_only: bool
+dispatch: Dispatch
+to_dict() dict
}
class Gates {
+writes_to_disk: bool
+requires_gpu: bool
+requires_serializable_input: bool
+sanitizes_output: bool
}
class IOSpec {
+data_keys: list[str]
+segment_data_keys: list[str]
+accepts: list[AudioForm]
+produces: list[ProducedForm]
}
class AudioToDocumentStage {
+BATCH_ONLY = True
+keep_keys: list[str] | None
+drop_keys: tuple[str, ...]
+serialize_segments: bool
+segments_key: str
+describe() StageContract
+_sanitize(data) dict
+_sanitize_nested(value) object
+process_batch(tasks) list
}
class SegmentExtractionStage {
+BATCH_ONLY = True
+output_dir: str
+output_key: str
+describe() StageContract
+process_batch(tasks) list
}
AgentReady <|-- AudioToDocumentStage
AgentReady <|-- SegmentExtractionStage
StageContract *-- Gates
StageContract *-- IOSpec
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def __init__( | ||
| self, | ||
| batch_size: int = 64, | ||
| keep_keys: list[str] | None = None, | ||
| drop_keys: tuple[str, ...] = (), | ||
| serialize_segments: bool = False, | ||
| segments_key: str = "segments", | ||
| ) -> None: | ||
| self.batch_size = batch_size | ||
| self.keep_keys = keep_keys | ||
| self.drop_keys = drop_keys | ||
| self.serialize_segments = serialize_segments | ||
| self.segments_key = segments_key |
There was a problem hiding this comment.
self.keep_keys is a list[str], so k not in keys is an O(n) scan on every iteration through data.items(). For a data dict with many keys and a non-trivial keep_keys whitelist, this becomes O(data_keys × keep_keys) per call to _sanitize. Since this runs for every task in every batch, converting to a frozenset at construction time is worth the one-time cost.
| def __init__( | |
| self, | |
| batch_size: int = 64, | |
| keep_keys: list[str] | None = None, | |
| drop_keys: tuple[str, ...] = (), | |
| serialize_segments: bool = False, | |
| segments_key: str = "segments", | |
| ) -> None: | |
| self.batch_size = batch_size | |
| self.keep_keys = keep_keys | |
| self.drop_keys = drop_keys | |
| self.serialize_segments = serialize_segments | |
| self.segments_key = segments_key | |
| def __init__( | |
| self, | |
| batch_size: int = 64, | |
| keep_keys: list[str] | None = None, | |
| drop_keys: tuple[str, ...] = (), | |
| serialize_segments: bool = False, | |
| segments_key: str = "segments", | |
| ) -> None: | |
| self.batch_size = batch_size | |
| self.keep_keys = frozenset(keep_keys) if keep_keys is not None else None | |
| self.drop_keys = frozenset(drop_keys) | |
| self.serialize_segments = serialize_segments | |
| self.segments_key = segments_key |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads=IOSpec(data_keys=[]), | ||
| writes=IOSpec(data_keys=[]), | ||
| cardinality="N:1", | ||
| # Strips tensors/audio blobs while building the DataFrame, so its | ||
| # output is serialization-safe — the sanctioned sink to place before | ||
| # a JSON writer when a resident tensor may be present. | ||
| gates=Gates(sanitizes_output=True), | ||
| description="Aggregate AudioTasks into a DocumentBatch, stripping tensors/audio blobs (JSON/disk-safe).", | ||
| ) |
There was a problem hiding this comment.
batch_only not reflected in describe() output
BATCH_ONLY = True is set on the class, but the StageContract returned by describe() keeps batch_only=False (the default). The doc comment on StageContract.batch_only says it is "auto-derived at discovery time", so the registry path is correct — but any caller that invokes stage.describe() directly (e.g. a test, a planner that bypasses the registry, or a stage that calls describe() on a sub-stage) receives a contract that says process() is safe to call when it actually raises NotImplementedError. Consider passing batch_only=True here explicitly; the same gap exists in SegmentExtractionStage.describe() and PreserveByValueStage.describe().
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | ||
| if k == self.segments_key and self.serialize_segments: | ||
| cleaned[k] = self._sanitize_nested(v) | ||
| continue |
There was a problem hiding this comment.
_DROP_VALUE sentinel returned by _sanitize_nested is stored into cleaned without a guard check. If segments_key's value happens to be a raw torch.Tensor (rather than the expected list-of-dicts), _sanitize_nested returns _DROP_VALUE and cleaned[k] becomes an unserializable Python sentinel object, which pandas will store as an object-dtype cell. Any downstream JSON or Parquet writer will then fail — silently violating the sanitizes_output=True contract declared in describe().
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | |
| if k == self.segments_key and self.serialize_segments: | |
| cleaned[k] = self._sanitize_nested(v) | |
| continue | |
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | |
| if k == self.segments_key and self.serialize_segments: | |
| nested = self._sanitize_nested(v) | |
| if nested is not _DROP_VALUE: | |
| cleaned[k] = nested | |
| continue |
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads_one_of=[ | ||
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | ||
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | ||
| ], | ||
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | ||
| gates=Gates(writes_to_disk=True), | ||
| ) |
There was a problem hiding this comment.
reads_one_of contract only covers Combo 2 (timestamp-based) and Combo 3 (diarization with diar_segments + speaker_id), but misses Combo 4 (speaker + timestamps: speaker_id present, diar_segments absent). An agent planning around this contract would incorrectly infer that no speaker_id input is needed for the timestamp combo, or that diar_segments is always required alongside speaker_id.
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads_one_of=[ | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | |
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | |
| gates=Gates(writes_to_disk=True), | |
| ) | |
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads_one_of=[ | |
| # Combo 2: timestamp-based extraction (no speaker info) | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | |
| # Combo 3: diarization-based extraction (diar_segments + speaker_id) | |
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | |
| # Combo 4: speaker + timestamp extraction (speaker_id present, no diar_segments) | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms", "speaker_id"], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | |
| gates=Gates(writes_to_disk=True), | |
| ) |
…pts/reads helpers)
1 participant
Summary
Makes the audio IO stages agent-ready (they now implement
describe()) and hardens the document sink.io/convert.py(AudioToDocumentStage) — N:1 sink, AudioTask → DocumentBatchkeep_keys(whitelist of columns to emit;None= all, prior behavior),drop_keys(blacklist),serialize_segments+segments_key.segments(a list of per-segment dicts that may embed waveform tensors) was previously dropped wholesale. Withserialize_segments=Trueit is kept but recursively cleaned by the new_sanitize_nestedhelper, which walks nested dicts/lists and stripstorch.Tensor/ audio-blob values (using a_DROP_VALUEsentinel so a legitimateNoneis preserved).describe()setsGates(sanitizes_output=True)to mark this as the tensor-safe boundary before a JSON/parquet writer.io/extract_segments.py(SegmentExtractionStage)describe()with OR-shaped reads covering the extraction combos (timestamp-based vs diarization-based). MarkedBATCH_ONLY: extraction runs per batch and appends to a shared metadata file, soprocess()raises andprocess_batch()does the work.Notes for reviewers
keep_keys=Noneemits all keys as before;segmentsstill dropped unlessserialize_segments=True._sanitize/_sanitize_nestedinconvert.py(recursive tensor stripping) and the combo auto-detection inextract_segments.py.